#########################################################################################
# 作者：hehung
# 联系：1398660197@qq.com
# 微信：hehung95
# 描述：该脚本为手势控制机械臂程序，使用手势控制即可，使用之前可能需要做标定，并不一定满足
#       所有人
#########################################################################################

import cv2
import numpy as np
import mediapipe as mp
import math
import uart_ctrl


g_wrist_x_angle = 0.0
g_wrist_y_angle = 0.0
g_hand_dist = 0.0
g_hand_x_norm = 0.0
g_hand_y_norm = 0.0
g_catch = 0.0
g_steer = [7.5, 7.5, 7.5, 7.5, 7.5, 7.5]
g_hand_dist_err = 0

g_steer_filter = [[7.5] * 10 for _ in range(6)]

FILTER_NUM = 10

# 图片处理类
class ImgClass:
    def __init__(self):
        # 获取摄像机
        self.cap = cv2.VideoCapture(0)

    def read_cap(self):
        # 读取摄像头采集的视频
        ret, frame = self.cap.read()
        return frame

    # 显示图片
    def show_img(self, img):
        cv2.imshow('Gesture Control Robotic Arm', img)
        c = cv2.waitKey(1)
        # 如果按下ESC则退出
        if c == 27:
            self.close_img()
            return 0
        else:
            return 1

    # 关闭显示
    def close_img(self):
        self.cap.release()
        cv2.destroyAllWindows()

class class_mediapipe():
    def __init__(self):
        self.uart_app = uart_ctrl.class_uart()
        # 手势识别
        self.mp_hands = mp.solutions.hands
        self.hands = self.mp_hands.Hands()
        self.mp_drawing = mp.solutions.drawing_utils

    # 滤波
    def average_filter(self, steer_num, value):
        for i in range(FILTER_NUM):
            if (i < FILTER_NUM - 1):
                g_steer_filter[steer_num][i] = g_steer_filter[steer_num][i+1]
            else:
                g_steer_filter[steer_num][i] = value
            g_steer[steer_num] = g_steer[steer_num] + g_steer_filter[steer_num][i]
        g_steer[steer_num] = g_steer[steer_num] / FILTER_NUM

    # 计算手指到手腕的距离
    def calc_finger_dist(self, landmarks, finger_tip_index1, finger_tip_index2):
        x1, y1 = landmarks.landmark[finger_tip_index1].x, landmarks.landmark[finger_tip_index1].y
        x2, y2 = landmarks.landmark[finger_tip_index2].x, landmarks.landmark[finger_tip_index2].y
        
        dist = math.hypot(x2 - x1, y2 - y1)
        return dist

    # 手势处理
    def mediapipe_process(self, frame):
        global g_wrist_x_angle
        global g_wrist_y_angle
        global g_hand_dist
        global g_hand_x_norm
        global g_hand_y_norm
        global g_catch
        global g_hand_dist_err

        # 将BGR转换为RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 手势识别处理
        results = self.hands.process(image)
        # 再将RGB转回BGR
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        # 判断是否检测到手势
        if results.multi_hand_landmarks:
            # 两只手都存在时，只处理一只手
            hand_landmarks = results.multi_hand_landmarks[0]

            ########################################################
            # 计算thumb_cmc到手腕的直线距离，用于计算手掌和摄像头之间的距离
            # 0 - 手腕
            # 1 - thumb_cmc
            g_hand_dist = self.calc_finger_dist(hand_landmarks, 0, 1)
            # 限定距离范围
            # 该参数值是标定的结果，实际使用需要根据摄像头不同以及放置位置区别进行实际标定
            g_hand_dist_err = 0
            if (g_hand_dist <= 0.04):
                # 距离太远
                g_hand_dist_err = 1
                g_hand_dist = 0.04
            elif (g_hand_dist >= 0.08):
                # 距离太近
                g_hand_dist_err = 2
                g_hand_dist = 0.08
            g_hand_dist = (g_hand_dist - 0.08) / (0.04 - 0.08)
            if (g_hand_dist < 0):
                g_hand_dist = 0
            elif (g_hand_dist > 1):
                g_hand_dist = 1
            # print (f'distance: {g_hand_dist}')
            
            ########################################################
            # 手在采集画面中的位置，通过手腕的左边要标记手掌在采集的视频图像中的位置
            # 手腕坐标
            wrist = hand_landmarks.landmark[0]
            # 保留X轴多表五位小数点
            g_hand_x_norm = np.round(wrist.x, 5)
            # 限定坐标范围
            if (g_hand_x_norm > 1):
                g_hand_x_norm = 1
            elif (g_hand_x_norm < 0):
                g_hand_x_norm = 0
            # 保留Y轴坐标五位小数
            g_hand_y_norm = np.round(wrist.y, 5)
            # 限定坐标范围
            if (g_hand_y_norm < 0.4):
                g_hand_y_norm = 0.4
            elif (g_hand_y_norm > 1):
                g_hand_y_norm = 1
            g_hand_y_norm = (g_hand_y_norm - 0.4) * 1.6
            
            # print ('Hand position:' + str(g_hand_x_norm) + '     ' + str(g_hand_y_norm))

            ########################################################
            # 计算中指到大拇指的距离（用于计算手掌是否握拳以及握拳的弯曲程度）
            # 0 - 手腕wist
            # 12 - 中指指尖middle_finger_tip
            middle_dist = np.round(self.calc_finger_dist(hand_landmarks, 0, 12), 5)
            # 理论距离公式（手掌伸开时）
            middle_dist_norm_extend = (-0.24 * g_hand_dist + 0.44)
            # 理论距离公式（手掌握拳时）
            middle_dist_norm_fist   = (-0.08 * g_hand_dist + 0.2)
            # 差值
            if (middle_dist > middle_dist_norm_extend):
                middle_dist = middle_dist_norm_extend
            elif (middle_dist < middle_dist_norm_fist):
                middle_dist = middle_dist_norm_fist
            # 归一化 0 ~ 1
            fist_v_per = (middle_dist - middle_dist_norm_fist) / (middle_dist_norm_extend - middle_dist_norm_fist)
            g_catch = 1 - fist_v_per

            # print ("Finger distances:", middle_dist, g_hand_dist, g_catch)

            ########################################################
            # 计算手掌绕手腕翻转角度
            # 3 - 大拇指关节thumb_ip
            # 17 - 小拇指关节pinky_mcp
            thumb_mcp_x, thumb_mcp_y = hand_landmarks.landmark[3].x * 300, hand_landmarks.landmark[3].y * 400
            pinky_mcp_x, pinky_mcp_y = hand_landmarks.landmark[17].x * 300, hand_landmarks.landmark[17].y * 400    
            # 计算角度
            dx = pinky_mcp_x - thumb_mcp_x  
            dy = pinky_mcp_y - thumb_mcp_y
            g_wrist_x_angle = math.atan2(dy, dx) * 180 / math.pi
            if (g_wrist_x_angle < -80):
                g_wrist_x_angle = -80
            elif (g_wrist_x_angle > 80):
                g_wrist_x_angle = 80
            g_wrist_x_angle = g_wrist_x_angle + 80
            g_wrist_x_angle = 160 - g_wrist_x_angle
            # range 0 ~ 160
            # print(f"Palm rotation angle: {g_wrist_x_angle} degrees")

            self.mp_drawing.draw_landmarks(
                image,
                hand_landmarks,
                self.mp_hands.HAND_CONNECTIONS) # 用于指定地标如何在图中连接。
        else:
            g_hand_dist_err = 3
            for i in range(6):
                g_steer[i] = 7.5

        return image

    def pose_calc(self):
        global g_wrist_x_angle
        global g_wrist_y_angle
        global g_hand_dist
        global g_hand_x_norm
        global g_hand_y_norm
        global g_catch
        global g_steer

        # 控制逻辑为0~100%占空比，100占空比对应的20ms
        # SG90舵机PWM控制： 0度周期0.5ms，对应占空比为2.5%
        #                  180度周期2.5ms，对应占空比12.5%
        # 旋转计算
        # 控制舵机旋转，中间为90度，对应SG90舵机PWM控制占空比7.5%
        # 舵机0度到180度，对应2.5%~12.5%
        g_hand_x_norm = 1 - g_hand_x_norm
        g_steer[0] = 2.5 + (g_hand_x_norm) * 10
        # 机械臂夹子夹的力度计算
        # 取值为0 ~ 0.2，0对应没握拳，0.2对应全握
        # 计算公式如下
        g_steer[5] = 7.5 + (g_catch) * 5
        # 翻转, 范围：0~100
        g_steer[4] = 2.5 + (g_wrist_x_angle) / 18
        # 上下
        # 取值范围 0~1
        g_hand_dist = 1 - g_hand_dist
        g_steer[1] = 7.5 + (g_hand_dist) * 5

        # 前后
        g_steer[2] = 7.5 - (g_hand_y_norm) * 5
        g_steer[3] = 7.5 - (g_hand_y_norm) * 5

        # average_filter(0, g_steer[0])
        # average_filter(1, g_steer[1])
        # average_filter(2, g_steer[2])
        # average_filter(4, g_steer[4])
        # average_filter(5, g_steer[5])

        # for i in range(6):
        #     # 取小数点位数为2
        #     g_steer[i] = round(g_steer[i], 2)
        #     print (g_steer[i], end = ' -- ')
        # print ()

    # 在图片上添加文字信息
    def image_add_text(self, img):
        global g_hand_dist_err

        font = cv2.FONT_HERSHEY_SIMPLEX 
        font_scale = 1
        thickness = 2

        # 在图像上绘制文本
        text = 'Gesture Control Robotic Arm'
        # 文本起始坐标
        org = (250, 30)
        # 蓝色文字
        color = (255, 0, 0)
        cv2.putText(img, text, org, font, font_scale, color, thickness)
        text = 'by hehung'
        # 文字坐标
        org = (600, 60)
        cv2.putText(img, text, org, font, font_scale, color, thickness)

        if g_hand_dist_err != 0:
            if g_hand_dist_err == 1:
                text = 'The hand is too close in distance'
            elif g_hand_dist_err == 2:
                text = 'The hand is too far away in distance'
            elif g_hand_dist_err == 3:
                text = '             No Hand Found'
            # 红色文字
            color = (0, 0, 255)
            # 文字坐标
            org = (220, 780)
            # 显示文字
            cv2.putText(img, text, org, font, font_scale, color, thickness)

        return img

    # 串口数据发送
    def send_command_to_servo(self):
        # data_str = 'START::0-%.2f:1-%.2f:2-%.2f:3-%.2f:4-%.2f:5-%.2f::END\r\n' % (g_steer[0], g_steer[1], g_steer[2], g_steer[3], g_steer[4], g_steer[5])
        # print ('Send Data: %s' % data_str.encode())
        # self.uart_app.send(data_str)
        self.uart_app.servo_ctrl(g_steer)

    def gesture_main(self):
        img_handle = ImgClass()
        while True:
            # 采集摄像机图像并返回图片数据
            frame = img_handle.read_cap()
            # 采集的图像镜像，便于直观查看
            frame = cv2.flip(frame, 1)
            # 手势处理
            img = self.mediapipe_process(frame)
            # 计算每个舵机的控制命令
            self.pose_calc()
            # 将控制命令通过串口发送到舵机控制器
            self.send_command_to_servo()
            # 放大图片，便于显示屏查看
            img = cv2.resize(img, (1080, 810)) 
            # 在图片上添加版权信息
            img = self.image_add_text(img)
            # 显示图片
            if (0 == img_handle.show_img(img)):
                break

if __name__ == '__main__':
    pass
